-------------------------------------------------------------------------------
-- Copas - Coroutine Oriented Portable Asynchronous Services
-- 
-- Offers a dispatcher and socket operations based on coroutines.
-- Usage:
--    copas.addserver(server, handler, timeout)
--    copas.addthread(thread, ...) Create a new coroutine thread and run it with args
--    copas.loop(timeout) - listens infinetely
--    copas.step(timeout) - executes one listening step
--    copas.receive(pattern or number) - receives data from a socket
--    copas.settimeout(client, time) if time=0 copas.receive(bufferSize) - receives partial data from a socket were data<=bufferSize
--    copas.send  - sends data through a socket
--    copas.wrap  - wraps a LuaSocket socket with Copas methods
--    copas.connect - blocks only the thread until connection completes
--    copas.flush - *deprecated* do nothing
--
-- Authors: Andre Carregal and Javier Guerra
-- Contributors: Diego Nehab, Mike Pall, David Burgess, Leonardo Godinho
--               and Thomas Harning Jr.
--
-- Copyright 2007 - Kepler Project (www.keplerproject.org)
--
-- $Id: copas.lua,v 1.26 2007/06/20 21:57:16 carregal Exp $
-------------------------------------------------------------------------------
local socket = require "socket"
WATCH_DOG_TIMEOUT = 120
local time = os.time
local duration = os.difftime

SELECT = socket.select

local function proto_index(o, k)  
    local v = o.__proto[k]
    if type(v) == "function" then return function(x,...) return v(o.__proto,...) end 
    else return v end
end

function prototype(o)
    return setmetatable({__proto=o}, {__index = proto_index })
end

-- Redefines LuaSocket functions with coroutine safe versions
-- (this allows the use of socket.http from within copas)
local function statusHandler(status, ...)
  if status then return ... end
  return nil, ...
end
function socket.protect(func)
  return function (...)
		return statusHandler(pcall(func, ...))
	end
end

function socket.newtry(finalizer)
	return function (...)
		local status = (...) or false
		if (status==false)then
			pcall(finalizer, select(2, ...) )
			error((select(2, ...)), 0)
		end
		return ...
	end
end
-- end of LuaSocket redefinitions


module ("copas3", package.seeall)

-- Meta information is public even if begining with an "_"
_COPYRIGHT   = "Copyright (C) 2004-2007 Kepler Project"
_DESCRIPTION = "Coroutine Oriented Portable Asynchronous Services"
_VERSION     = "Copas 1.1.1"

-------------------------------------------------------------------------------
-- Simple set implementation based on LuaSocket's tinyirc.lua example
-- adds a FIFO queue for each value in the set
-------------------------------------------------------------------------------
local function newset()
    local reverse = {}
    local set = {}
    local q = {}
    setmetatable(set, { __index = {
        insert = function(set, value)
            if not reverse[value] then
                set[#set + 1] = value
                reverse[value] = #set
            end
        end,

        has = function(set, value)
            return reverse[value]
        end,

        remove = function(set, value)
            local index = reverse[value]
            if index then
                reverse[value] = nil
                local top = set[#set]
                set[#set] = nil
                if top ~= value then
                    reverse[top] = index
                    set[index] = top
                end
            end
        end,
		
		push = function (set, key, itm)
			local qKey = q[key]
			if qKey == nil then
				q[key] = {itm}
			else
				qKey[#qKey + 1] = itm
			end
		end,
		
        pop = function (set, key)
          local t = q[key]
          if t ~= nil then
            local ret = table.remove (t, 1)
            if t[1] == nil then
              q[key] = nil
            end
            return ret
          end
        end,
    }})
    return set
end

local _servers = newset() -- servers being handled

_reading = newset() -- sockets currently being read
_writing = newset() -- sockets currently being written
_killing = {}
-------------------------------------------------------------------------------
-- Coroutine based socket I/O functions.
-------------------------------------------------------------------------------
-- reads a pattern from a client and yields to the reading set on timeouts
function receive(client, pattern, part)
  local suspended = 0
  local s, err
  pattern = pattern or "*l"
  if type(client) == "userdata" then client = prototype(client) end
  repeat
    local now = time()
    s, err, part = client:receive(pattern, part)
    if s or err ~= "timeout" then 
        client._copas_last_timeout = nil
        client._copas_reading = nil
        return s, err, part 
    end
    suspended = suspended + 1 
    assert(suspended < 1000)
    local probe = client.probe
    local last_timeout = client._copas_last_timeout
    if type(probe)=="function" and 
        last_timeout and 
        duration(now, last_timeout) > 2*WATCH_DOG_TIMEOUT then
        probe(client)
    end
    if not last_timeout then client._copas_last_timeout = now end
    client._copas_reading = now
    coroutine.yield(client, _reading)
  until false
end

-- same as above but with special treatment when reading chunks,
-- unblocks on any data received.
function receivePartial(client, pattern)
  local s, err, part
  pattern = pattern or "*l"
  if type(client) == "userdata" then client = prototype(client) end
  repeat
    local now = time()
    s, err, part = client:receive(pattern)
    if s or ( (type(pattern)=="number") and part~="" and part ~=nil ) or 
       err ~= "timeout" then 
        client._copas_last_timeout = nil
        client._copas_reading = nil
        return s, err, part 
    end
    local probe = client.probe
    local last_timeout = client._copas_last_timeout
    if type(probe)=="function" and 
        last_timeout and 
        duration(now, last_timeout) > 2*WATCH_DOG_TIMEOUT then
        probe(client)
    end
    if not last_timeout then client._copas_last_timeout = now end
    client._copas_reading = now
    coroutine.yield(client, _reading)
  until false
end

-- sends data to a client. The operation is buffered and
-- yields to the writing set on timeouts
function send(client,data, from, to)
  local s, err,sent
  from = from or 1
  local lastIndex = from - 1

  if type(client) == "userdata" then client = prototype(client) end
  repeat
    s, err, lastIndex = client:send(data, lastIndex + 1, to)
    -- adds extra corrotine swap
    -- garantees that high throuput dont take other threads to starvation 
    if (math.random(100) > 90) then
        client._copas_writing = time()
    	coroutine.yield(client, _writing)
    end
    if s or err ~= "timeout" then 
        client._copas_writing = nil
        return s, err,lastIndex 
    end
    client._copas_writing = time()
    coroutine.yield(client, _writing)
  until false
end

-- waits until connection is completed
function connect(skt,host, port)
	skt:settimeout(0)
  repeat
    local ret,err = skt:connect (host, port)
    if ret or err ~= "timeout" then 
      if (err=="already connected") then
        return 1
      end
      return ret, err 
    end
    if type(skt) == "userdata" then skt = prototype(skt) end
    skt._copas_writing = time()
    coroutine.yield(skt, _writing)
    skt._copas_writing = nil
  until false
end

-- flushes a client write buffer (deprecated)
function flush(client)
end

-- wraps a socket to use Copas methods (send, receive, flush and settimeout)
local _skt_mt = {__index = {
	send = function (self, data, from, to)
			return send (self.socket, data, from, to)
	end,
	
	receive = function (self, pattern)
			if (self.timeout==0) then
  				return receivePartial(self.socket, pattern)
  			end
			return receive (self.socket, pattern)
	end,
	
	flush = function (self)
			return flush (self.socket)
	end,
	
	settimeout = function (self,time)
			self.timeout=time
			return
	end,
}}

function wrap (skt)
	return  setmetatable ({socket = skt}, _skt_mt)
end

--------------------------------------------------
-- Error handling
--------------------------------------------------

local _errhandlers = {}   -- error handler per coroutine

function setErrorHandler (err)
	local co = coroutine.running()
	if co then
		_errhandlers [co] = err
	end
end

local function _deferror (msg, co, skt)
	print (msg, co, skt)
end

-------------------------------------------------------------------------------
-- Thread handling
-------------------------------------------------------------------------------

local function _doTick (co, skt, ...)
	if not co then return end
	
	local ok, res, new_q = coroutine.resume(co, skt, ...)
	
	if ok and res and new_q then
		new_q:insert (res)
		new_q:push (res, co)
	else
		if not ok then pcall (_errhandlers [co] or _deferror, res, co, skt) end
		if skt then skt:close() end
		_errhandlers [co] = nil
	end
end

-- accepts a connection on socket input
local function _accept(input, handler)
	local client = input:accept()
	if client then 
		client:settimeout(0)
		local co = coroutine.create(handler)
		_doTick (co, client)
		--_reading:insert(client)
	end
	return client
end

-- handle threads on a queue
local function _tickRead (skt)
	_doTick (_reading:pop (skt), skt)
end

local function _tickWrite (skt)
	_doTick (_writing:pop (skt), skt)
end

-------------------------------------------------------------------------------
-- Adds a server/handler pair to Copas dispatcher
-------------------------------------------------------------------------------
function addserver(server, handler, timeout)
  server:settimeout(timeout or 0.1)
  _servers[server] = handler
  _reading:insert(server)
end

-------------------------------------------------------------------------------
-- Adds an new courotine thread to Copas dispatcher
-------------------------------------------------------------------------------
function addthread(thread, ...)
	local co = coroutine.create(thread)
	_doTick (co, nil, ...)
end

-------------------------------------------------------------------------------
-- tasks registering
-------------------------------------------------------------------------------

local _tasks = {}

local function addtaskRead (tsk)
	-- lets tasks call the default _tick()
	tsk.def_tick = _tickRead
	
	_tasks [tsk] = true
end

local function addtaskWrite (tsk)
	-- lets tasks call the default _tick()
	tsk.def_tick = _tickWrite
	
	_tasks [tsk] = true
end

local function tasks ()
	return next, _tasks
end

-------------------------------------------------------------------------------
-- main tasks: manage readable and writable socket sets
-------------------------------------------------------------------------------
-- a task to check ready to read events
local _readable_t = {
  events = function(self)
  	local i = 0
  	return function ()
  		i = i + 1
  		return self._evs [i]
  	end
  end,
  
  tick = function (self, input)
  	local handler = _servers[input]
  	if handler then
  		input = _accept(input, handler)
  	else
  		_reading:remove (input)
  		self.def_tick (input)
  	end
  end
}

addtaskRead (_readable_t)


-- a task to check ready to write events
local _writable_t = {
  events = function (self)
  	local i = 0
  	return function ()
  		i = i+1
  		return self._evs [i]
  	end
  end,
  
  tick = function (self, output)
  	_writing:remove (output)
  	self.def_tick (output)
  end
}

addtaskWrite (_writable_t)

local last_cleansing = 0
local next_select = SELECT
local function _select (timeout)
	local err
    local readable={}
    local writable={}
    local r={}
    local w={}
    local now = os.time()

	_readable_t._evs, _writable_t._evs, err = next_select(_reading, _writing, timeout)
	local r_evs, w_evs = _readable_t._evs, _writable_t._evs 
    if #r_evs + #w_evs == 0 then 
        next_select = socket.select
    else 
        next_select = SELECT 
    end
    if duration(now, last_cleansing) > WATCH_DOG_TIMEOUT then
        last_cleansing = now
        for k,skt in pairs(_reading) do
            local v = skt._copas_reading
            if v and not r_evs[skt] and duration(now, v) > WATCH_DOG_TIMEOUT then
                skt._copas_reading = nil
                r_evs[#r_evs + 1] = skt
                r_evs[skt] = #r_evs
            end
        end

        for k,skt in pairs(_writing) do
            local v = skt._copas_writing
            if v and not w_evs[skt] and duration(now, v) > WATCH_DOG_TIMEOUT then
                skt._copas_writing = nil
                w_evs[#w_evs + 1] = skt
                w_evs[skt] = #w_evs
            end
        end
    end

    if err == "timeout" and #r_evs + #w_evs > 0 then return nil
    else return err end
end


-------------------------------------------------------------------------------
-- Dispatcher loop step.
-- Listen to client requests and handles them
-------------------------------------------------------------------------------
function step(timeout)
  local err = _select (timeout)
  if err == "timeout" then return end

	if err then
		error(err)
	end
		
	for tsk in tasks() do
		for ev in tsk:events () do
			tsk:tick (ev)
		end
	end
end

function event_loop(timeout)
    local lasync = require 'lasync'
    local handle = lasync.handle()
    local timer = handle:event(-1,{}, 0):open(0.1)
    local l={}
    for k,v in pairs(_servers) do 
        l[#l+1] = k
    end
    local function check_accept()
        local r,w=socket.select(l,{},0)
        for k,v in ipairs(r) do 
           _readable_t:tick(v) 
        end
    end

    local dispatcher = {
        ["readable"] = function(self, ev)
            local skt = ev:object()
            if skt then
                local now=os.time()
                check_accept()
                if not _reading:has(skt) and 
                    _writing:has(skt) and 
                    skt._copas_writing and
                    duration(now,skt._copas_writing) > 1*WATCH_DOG_TIMEOUT then
                    _writable_t:tick(skt)
                else
                    _readable_t:tick(skt)
                end
                if skt and duration(now, skt._copas_last_timeout or now) > 1*WATCH_DOG_TIMEOUT then
                    print("readable but nothing read for over %i", duration(now,skt._copas_last_timeout))
                end
            end
        end,
        ["writable"] = function(self, ev)
            local skt = ev:object()
            if skt then
                check_accept()
                if not _writing:has(skt) and _reading:has(skt) then
                    local last_timeout = skt._copas_last_timeout
                    local now=os.time()
                    if last_timeout and duration(now, last_timeout) > 2*WATCH_DOG_TIMEOUT then
                        print("writable but is waiting for read")
                        print(lasync.probe(skt:getfd()))
                    end
                end
                if true or _writing:has(skt) then
                    _writable_t:tick(skt)
                end
            end
        end,
        ["timeout"] = function(self, ev)
            local skt = ev:object()
            check_accept()

            if _reading:has(skt) then 
                _readable_t:tick(skt) 
                return
            end

            for k,v in pairs(_writing) do
                if not v.dangling then v.dangling = os.time() 
                elseif os.difftime(os.time(), v.dangling) > 1200 then
                    print(os.time(), "writing", v, os.difftime(os.time(), v.dangling))
                    local c,e = lasync.probe(v:getfd())
                    if e then _writable_t:tick(v) end
                end
            end
            for k,v in pairs(_reading) do
                if type(v) ~= "userdata" then
                    if not v.dangling then v.dangling = os.time() 
                    elseif os.difftime(os.time(), v.dangling) > 1200 then
                        print(os.time(), "reading", v, os.difftime(os.time(), v.dangling))
                        local c,e = lasync.probe(v:getfd())
                        if e then _readable_t:tick(v) end
                    end
                end
            end
        end,
        }
    handle:loop(dispatcher)
end

-------------------------------------------------------------------------------
-- Dispatcher endless loop.
-- Listen to client requests and handles them forever
-------------------------------------------------------------------------------
function loop(timeout)
	while true do
		step(timeout)
	end
end
